/*
 * Gongdao.com Inc.
 * Copyright (c) 2020-2034 All Rights Reserved.
 */
package com.github.novicezk.midjourney.controller;

import com.alibaba.fastjson.JSONObject;
import com.github.novicezk.midjourney.ProxyProperties;
import com.github.novicezk.midjourney.service.TranslateService;
import io.swagger.annotations.Api;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.IOException;
import java.security.SecureRandom;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author 瑜登
 * @version 2023年05月26日 9:25 PM
 */
@Api(tags = "任务查询")
@RestController
@RequestMapping("/v2")
@Slf4j
public class MjNewController {

	@Autowired
	private ProxyProperties proxyProperties;

	@Autowired
	private  TranslateService translateService;

	private static final Map<String, String> taskMap = new ConcurrentHashMap<>();

	@RequestMapping("/imagine")
	public String doWork(@org.springframework.web.bind.annotation.RequestBody JSONObject jsonObject) {

		String action = jsonObject.getString("action");
		if (StringUtils.isEmpty(action) || "generate".equals(action)){
			String prompt = jsonObject.getString("prompt");

			String promptEn;
			int paramStart = prompt.indexOf(" --");
			if (paramStart > 0) {
				promptEn = this.translateService.translateToEnglish(prompt.substring(0, paramStart)).trim() + prompt.substring(paramStart);
			} else {
				promptEn = this.translateService.translateToEnglish(prompt).trim();
			}

			jsonObject.put("prompt", promptEn);
		}

		jsonObject.put("callback_url","http://hua.itbook123.com/mj/v2/callback");

		try {
			return invokeAPI(jsonObject.toJSONString());
		} catch (IOException e) {
			log.error(e.getMessage(),e);
			return "服务异常";
		}
	}

	@RequestMapping("/callback")
	public String callBack(@org.springframework.web.bind.annotation.RequestBody JSONObject jsonObject){

		log.info("callback {}", jsonObject.toJSONString());

		String taskId = jsonObject.getString("task_id");
		taskMap.put(taskId,jsonObject.toJSONString());

		return "success";
	}

	@RequestMapping("/getTask")
	public String getTask(@RequestParam("taskId") String taskId){
		return taskMap.getOrDefault(taskId,"");
	}

	private String invokeAPI(String json) throws IOException {
//		OkHttpClient client = new OkHttpClient();

		OkHttpClient client = getUnsafeOkHttpClient();

// 创建请求体
		MediaType mediaType = MediaType.parse("application/json; charset=utf-8");

//		String json = "{\"id\":123,\"name\":\"test\"}";

		RequestBody requestBody = RequestBody.create(mediaType, json);

		// 创建请求
		HttpUrl.Builder urlBuilder = HttpUrl.parse("https://api.zhishuyun.com/midjourney/imagine")
				.newBuilder();
		urlBuilder.addQueryParameter("token", proxyProperties.getZhishuToken());
		String url = urlBuilder.build().toString();

// 创建请求
		Request request = new Request.Builder()
				.url(url)
				.post(requestBody)
				.addHeader("accept", "application/json")
				.build();

// 发送请求并处理响应
		Response response = client.newCall(request).execute();
		String responseBody = response.body().string();

		return responseBody;
	}


	private static OkHttpClient getUnsafeOkHttpClient() {
		try {
			TrustManager[] trustAllCerts = new TrustManager[] {
					new X509TrustManager() {
						public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {}
						public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {}
						public X509Certificate[] getAcceptedIssuers() { return new X509Certificate[0]; }
					}
			};

			SSLContext sslContext = SSLContext.getInstance("SSL");
			sslContext.init(null, trustAllCerts, new SecureRandom());

			return new OkHttpClient.Builder()
					.sslSocketFactory(sslContext.getSocketFactory(), (X509TrustManager) trustAllCerts[0])
					.hostnameVerifier((hostname, session) -> true)
					.build();
		} catch (Exception e) {
			throw new RuntimeException(e);
		}
	}


}
